from torchvision import models
import torchvision.transforms as transforms
from sklearn.metrics.pairwise import cosine_similarity

from PIL import Image, ImageFile
import torch


class Similarity:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
        self.model = torch.nn.Sequential(*list(self.model.children())[:-1])
        self.model.to(self.device).eval()

        self.meta_feature = None

        self.preprocess = transforms.Compose(
            [
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
                ),
            ]
        )

    def compare_features(self, feature1, feature2):
        similarity = cosine_similarity(feature1.reshape(1, -1), feature2.reshape(1, -1))
        return similarity[0][0]

    def get_feature(self, input_data: torch.Tensor):
        input_data = input_data.to(self.device)
        with torch.no_grad():  # 禁用梯度计算
            feature: torch.Tensor = self.model(
                input_data
            )  # 运行模型，hook会捕获中间层的特征
        return feature.cpu()

    def preprocess_image_file(self, image: ImageFile.ImageFile):
        image = self.preprocess(image.convert("RGB")).unsqueeze(0)  # 增加一个batch维度
        return image

    def get_similarity(self, data1, data2):
        return self.compare_features(self.get_feature(data1), self.get_feature(data2))

    def load_image(self, image_path):
        image = Image.open(image_path)
        return self.preprocess_image_file(image)

    def flask_interface(self, data1: ImageFile.ImageFile, data2: ImageFile.ImageFile):
        data1 = self.preprocess_image_file(data1)
        data2 = self.preprocess_image_file(data2)
        return float(self.get_similarity(data1, data2))

    def flask_interface_path(self, data_path1: str, data_path2: str):
        data1 = self.load_image(data_path1)
        data2 = self.load_image(data_path2)
        return float(self.get_similarity(data1, data2))


def test():
    similarity = Similarity()
    img1 = similarity.load_image(
        "/home/tuchunxu/workspace/java/pets_pic/siamese_cat1.jpeg"
    )
    img2 = similarity.load_image(
        "/home/tuchunxu/workspace/java/pets_pic/siamese_cat4.jpeg"
    )

    res = similarity.get_similarity(img1, img2)
    print(res)


if __name__ == "__main__":
    test()
